#!/usr/bin/env python3
from rdkit import Chem
from collections import Counter
from concurrent.futures import ProcessPoolExecutor
import pandas as pd
import argparse


def count_atoms_in_molecule(smiles):
    try:
        mol = Chem.MolFromSmiles(smiles)
        if mol is not None:
            return Counter(atom.GetSymbol() for atom in mol.GetAtoms())
        else:
            return Counter()
    except:
        return Counter()


def count_atoms_in_smiles_array(smiles_array, num_workers=None):
    with ProcessPoolExecutor(max_workers=num_workers) as executor:
        results = list(executor.map(count_atoms_in_molecule, smiles_array))
    return results


def extract_non_hydrogen_counts(atom_counts_list):
    non_hydrogen_counts = []
    for atom_counts in atom_counts_list:
        # Subtract the count of hydrogen atoms (if present) from the total count
        total_non_hydrogen = sum(
            count for element, count in atom_counts.items() if element != "H"
        )
        non_hydrogen_counts.append(total_non_hydrogen)
    return non_hydrogen_counts


def main():
    parser = argparse.ArgumentParser(
        description="Count non-hydrogen atoms in a BindingDB DTI dataset and filter by HAC"
    )

    parser.add_argument(
        "-I",
        "--input",
        type=str,
        required=True,
        help="Path to the BindingDB DTI data file (TSV format, can be gz compressed).",
    )
    parser.add_argument(
        "-M", "--max_hac", type=int, help="Maximum HAC to include in the dataset."
    )
    parser.add_argument(
        "-O",
        "--output",
        type=str,
        help="Path to the output file (TSV format). Required and only work when --max_hac is specified.",
    )

    args = parser.parse_args()

    path = args.input
    max_hac = args.max_hac
    out_path = args.output

    print("Reading dataset...")
    df = pd.read_csv(path, sep="\t")

    print("Counting atoms...")
    atom_counts = count_atoms_in_smiles_array(df["Ligand SMILES"])

    df["HAC"] = extract_non_hydrogen_counts(atom_counts)

    print(f"Max HAC in input drug library: {max(df['HAC'])}")
    print(f"Number of rows: {len(df)}")

    if max_hac is not None:
        print(f"Discard drugs with HAC > {max_hac} and problematic in rdkit process...")
        df = df[(df["HAC"] <= max_hac) & (df["HAC"] > 0)]
        print(f"Number of rows after filtering: {len(df)}")

        df.to_csv(out_path, index=False, sep="\t")


if __name__ == "__main__":
    main()
